import matplotlib.pyplot as plt
import networkx as nx  # 用于拓扑排序和图操作
import pygraphviz as pgv
from PySide6.QtCore import QPointF


class 有向图布局器:
    def __init__(self):
        self.图形 = pgv.AGraph(directed=True)  # 使用 Graphviz AGraph 表示有向图
        self.节点大小字典 = {}  # 记录每个节点的大小（像素）
        self.节点布局 = {}  # 记录每个节点的布局位置
        self.最小宽度 = 0  # 图形的最小宽度
        self.最小高度 = 0  # 图形的最小高度
        self.nx_图形 = nx.DiGraph()  # 使用 NetworkX 的图形来优化性能

    def _调整大小(self, 大小值):
        """根据规则调整节点大小为 100 的整数倍"""
        if 大小值 < 100:
            return 100
        else:
            # 计算需要的倍数
            倍数 = (大小值 + 99) // 100  # 向上取整到最近的 100 的倍数
            return 倍数 * 100

    def _像素转英寸(self, 宽度像素, 高度像素, dpi=72):
        """将像素大小转换为英寸大小"""
        宽度英寸 = 宽度像素 / dpi
        高度英寸 = 高度像素 / dpi
        return 宽度英寸, 高度英寸

    def 添加节点(self, 节点名称, 大小=(1.0, 1.0)):
        """添加节点并记录其大小，大小按照 100 的整数倍进行调整"""
        调整后宽度 = self._调整大小(大小[0])
        调整后高度 = self._调整大小(大小[1])

        # 将调整后的大小转换为适合 Graphviz 的尺寸（以英寸为单位，默认 72 DPI）
        宽度英寸, 高度英寸 = self._像素转英寸(调整后宽度, 调整后高度)

        self.图形.add_node(
            节点名称,
            width=宽度英寸,
            height=高度英寸,
            shape="box",
            fixedsize="true"
        )
        self.节点大小字典[节点名称] = (调整后宽度, 调整后高度)
        self.nx_图形.add_node(节点名称)

    def 修改节点大小(self, 节点名称, 新大小):
        """修改已有节点的大小，按照 100 的整数倍调整"""
        if 节点名称 in self.节点大小字典:
            调整后宽度 = self._调整大小(新大小[0])
            调整后高度 = self._调整大小(新大小[1])

            宽度英寸, 高度英寸 = self._像素转英寸(调整后宽度, 调整后高度)

            self.图形.get_node(节点名称).attr['width'] = 宽度英寸
            self.图形.get_node(节点名称).attr['height'] = 高度英寸
            self.节点大小字典[节点名称] = (调整后宽度, 调整后高度)
        else:
            print(f"节点 {节点名称} 不存在，无法修改大小")

    def 删除节点(self, 节点名称):
        """删除节点及其相关信息"""
        if self.图形.has_node(节点名称):
            self.图形.delete_node(节点名称)
            del self.节点大小字典[节点名称]
            self.nx_图形.remove_node(节点名称)

    def 添加边(self, 起始节点, 结束节点, 检查环路=True):
        """添加边并检查环路，避免重复转换图形"""
        if self.nx_图形.has_edge(起始节点, 结束节点):
            return  # 已存在边，避免重复添加

        self.nx_图形.add_edge(起始节点, 结束节点)
        if 检查环路:
            try:
                nx.find_cycle(self.nx_图形, orientation='original')
                # 如果找到环，移除边
                self.nx_图形.remove_edge(起始节点, 结束节点)
                print(f"添加边 {起始节点} -> {结束节点} 会导致环路，操作已取消")
            except nx.NetworkXNoCycle:
                # 无环，正式添加边到图形
                self.图形.add_edge(起始节点, 结束节点)
        else:
            self.图形.add_edge(起始节点, 结束节点)

    def 删除边(self, 起始节点, 结束节点):
        """删除边"""
        if self.图形.has_edge(起始节点, 结束节点):
            self.图形.delete_edge(起始节点, 结束节点)
            self.nx_图形.remove_edge(起始节点, 结束节点)

    def 计算拓扑排序(self):
        """使用 NetworkX 进行拓扑排序"""
        return list(nx.topological_sort(self.nx_图形))

    def 计算布局(self, 布局算法='dot', dpi=72, 布局方向="LR", x_缩放因子=88, y_缩放因子=55):
        """计算节点布局，消除空白空间并调整节点位置"""
        # 设置图形属性
        self.图形.graph_attr.update(
            dpi=f"{dpi}",
            overlap="false",
            splines="false",
            nodesep="0.5",
            ranksep="0.8",
            sep="+15",
            ratio="auto",
            margin="0.2,0.2",
            rankdir=布局方向
        )

        # 应用布局算法
        self.图形.layout(prog=布局算法)

        # 获取节点坐标并调整位置
        节点坐标 = {}  # 存储每个节点的坐标
        x坐标列表 = []  # 存储所有节点的 x 坐标
        y坐标列表 = []  # 存储所有节点的 y 坐标

        for 节点 in self.图形.nodes():
            位置属性 = self.图形.get_node(节点).attr.get("pos", None)
            if 位置属性:
                x字符串, y字符串 = 位置属性.split(",")
                x坐标 = float(x字符串)
                y坐标 = float(y字符串)
                节点坐标[节点] = (x坐标, y坐标)
                x坐标列表.append(x坐标)
                y坐标列表.append(y坐标)
            else:
                节点坐标[节点] = (0.0, 0.0)
                x坐标列表.append(0.0)
                y坐标列表.append(0.0)

        if x坐标列表 and y坐标列表:
            x最小值 = min(x坐标列表)
            y最小值 = min(y坐标列表)
            x最大值 = max(x坐标列表)
            y最大值 = max(y坐标列表)
        else:
            x最小值 = y最小值 = x最大值 = y最大值 = 0

        self.节点布局 = {}

        for 节点, (x坐标, y坐标) in 节点坐标.items():
            x调整后 = (x坐标 - x最小值)
            y调整后 = (y坐标 - y最小值)
            self.节点布局[节点] = QPointF(x调整后 * dpi / x_缩放因子, y调整后 * dpi / y_缩放因子)

        # 更新最小宽度和高度
        self.最小宽度 = (x最大值 - x最小值) * dpi / x_缩放因子
        self.最小高度 = (y最大值 - y最小值) * dpi / y_缩放因子

    def 获取孤立节点(self):
        """获取图中的孤立节点"""
        return list(nx.isolates(self.nx_图形))

    def 获取起始节点(self):
        """获取入度为 0 的起始节点"""
        return [节点 for 节点 in self.nx_图形.nodes() if self.nx_图形.in_degree(节点) == 0]

    def 获取分区(self):
        """获取弱连通分量（子图）"""
        return [list(c) for c in nx.weakly_connected_components(self.nx_图形)]

    def 获取所有子图的起始节点(self):
        """获取所有子图的起始节点"""
        分区列表 = self.获取分区()
        起始节点列表 = []
        for 子图节点集合 in 分区列表:
            子图 = self.nx_图形.subgraph(子图节点集合)
            起始节点 = [节点 for 节点 in 子图.nodes() if 子图.in_degree(节点) == 0]
            起始节点列表.extend(起始节点)
        return 起始节点列表

    def 获取节点位置(self, 节点名称):
        """获取指定节点的位置"""
        return self.节点布局.get(节点名称, None)

    def 获取所有节点位置(self):
        """获取所有节点的位置字典"""
        return self.节点布局

    def 绘制图形(self, 输出文件="output.png"):
        """绘制并显示图形"""
        self.图形.draw(输出文件, format="png")
        img = plt.imread(输出文件)
        plt.imshow(img)
        plt.axis("off")
        plt.show()
